import time

from tinyops.vendor import justqueue
from tinyops.config import Config, DirectoryManager

from .task import Task


class TaskQueue:
    """
    Task queue, don't share between threads and processors.
    """

    def __init__(self, config: Config):
        self.config = config
        self._queue = None

    def __enter__(self):
        self._init()
        return self

    def __exit__(self, *exc_info):
        self._close()

    def close(self):
        self._close()

    def _init(self):
        dm = DirectoryManager(self.config)
        filepath = dm.queue()
        self._queue = justqueue.FIFOQueue(filepath)

    def _close(self):
        if self._queue is not None:
            self._queue.close(remove=False)

    def _get_queue(self):
        if self._queue is None:
            self._init()
        return self._queue

    def submit(self, task_name: str, env: dict = None) -> Task:
        task = Task.from_config(task_name, self.config, env=env)
        self.put(task)
        return task

    def put(self, task: Task):
        """
        Put task into queue
        """
        self._get_queue().push(task.to_json())

    def get(self, timeout=None) -> Task:
        """
        Get task from queue

        Args:
            timeout (int):
                - None: wait until get result
                - 0: return result or None immediate
                - >0: wait until timeout, when timeout return None
        """
        item = self._poll_item(timeout=timeout)
        if item:
            return Task.from_json(item, self.config)
        return None

    def _poll_item(self, timeout=None):
        """
        use polling to support put and get between multiprocessor.
        """
        queue = self._get_queue()
        if timeout is None:
            while True:
                try:
                    return queue.pop()
                except justqueue.exceptions.EmptyQueueError:
                    time.sleep(0.1)
        if timeout <= 0:
            try:
                return queue.pop()
            except justqueue.exceptions.EmptyQueueError:
                return None
        deadline = time.time() + timeout
        while time.time() <= deadline:
            try:
                return queue.pop()
            except justqueue.exceptions.EmptyQueueError:
                time.sleep(0.1)
        return None
